import matplotlib.pyplot as plt  # 导入matplotlib.pyplot模块，用于绘图
from gzip import open as gzip_open  # 从gzip模块导入open函数，用于读取gzip压缩文件
from pickle import load as pickle_load  # 从pickle模块导入load函数，用于加载pickle文件
from pathlib import Path  # 从pathlib模块导入Path类，用于处理文件路径

from src.config.config import Config  # 从src.config.config模块导入Config类，用于获取配置信息
from src.config.config_plotting import PlotConfig, generic_styling  # 从src.config.config_plotting模块导入PlotConfig类和generic_styling函数

def plot_training_graph(
        path,  # 数据文件路径
        name,  # 图形名称
        xlabel: None or str = None,  # X轴标签，默认为None
        ylabel: None or str = None,  # Y轴标签，默认为None
):

    plot_config = PlotConfig()  # 创建PlotConfig类的实例，用于获取绘图配置信息

    width = plot_config.textwidth  # 设置图形宽度
    height = width * 2 / 3  # 设置图形高度，宽高比为2:3

    # 打开gzip压缩文件并加载数据
    with gzip_open(path, 'rb') as file:
        data = pickle_load(file)  # 使用pickle加载数据

    # 创建绘图窗口和坐标轴
    fig, ax = plt.subplots(figsize=(width, height))
    
    # 绘制散点图
    ax.scatter(range(len(data['mean_sum_rate_per_episode'])), data['mean_sum_rate_per_episode'])

    # 设置X轴和Y轴标签
    if xlabel:
        ax.set_xlabel(xlabel)
    if ylabel:
        ax.set_ylabel(ylabel)

    # 应用通用样式
    generic_styling(ax=ax)
    fig.tight_layout(pad=0)  # 自动调整布局

    # 保存图形
    plot_config.save_figures(plot_name=name, padding=0)


if __name__ == '__main__':  # 确保仅在直接运行脚本时执行以下代码
    cfg = Config()  # 创建Config类的实例，用于加载配置信息
    
    # 定义数据文件路径
    path = Path(cfg.output_metrics_path,
                'sat_2_ant_4_usr_3_satdist_10000_usrdist_1000',
                'err_mult_on_steering_cos',
                'single_error',
                'training_error_0.0_userwiggle_30.gzip')

    # 调用自定义绘图函数，生成训练图
    plot_training_graph(path, name='training_test', xlabel='Training Episode', ylabel='Mean Reward')
    plt.show()  # 显示绘制的图形
